import cv2
import os
import sys

# Width, Height
Camera_PROP_640x480 = (640, 480)
Camera_PROP_800x600 = (800, 600)
Camera_PROP_960_640 = (960, 640)
Camera_PROP_1024x540 = (1024, 540)
Camera_PROP_1024x600 = (1024, 600)
Camera_PROP_1024x768 = (1024, 768)
Camera_PROP_1080x960 = (1080, 960)
Camera_PROP_1152x864 = (1152, 864)
Camera_PROP_1280x600 = (1280, 600)
Camera_PROP_1280x720 = (1280, 720)
Camera_PROP_1280x768 = (1280, 768)
Camera_PROP_1280x800 = (1280, 800)
Camera_PROP_1280x960 = (1280, 960)
Camera_PROP_1280x1024 = (1280, 1024)
Camera_PROP_1360x768 = (1360, 768)
Camera_PROP_1366x768 = (1366, 768)
Camera_PROP_1400x1050 = (1400, 1050)
Camera_PROP_1440x900 = (1440, 900)
Camera_PROP_1600x900 = (1600, 900)
Camera_PROP_1680x1050 = (1680, 1050)
Camera_PROP_1920x1080 = (1920, 1080)
Camera_PROP_2048x1080 = (2048, 1080)
Camera_PROP_4096x2160 = (4096, 2160)
Camera_PROP_Max = (300000, 300000)

# 解决显示图片过大导致图片跑到显示器外边
def imshow(winname, mat):
    ratio = mat.shape[0] // 600
    if ratio <= 1:
        ratio = 1.5
    cv2.namedWindow(winname, cv2.WINDOW_NORMAL)
    cv2.resizeWindow(winname, int(mat.shape[1] / ratio), int(mat.shape[0] / ratio))
    cv2.imshow(winname, mat)

if __name__ == '__main__':
    print("=============================================")
    print("=  热键(请在摄像头的窗口使用）                   =")
    print("=  w: 拍摄图片                                =")
    print("=  q: 退出                                   =")
    print("=============================================")
    # 新建文件夹 test
    class_name = 'test'
    if not os.path.exists(class_name):
        os.mkdir(class_name)
    index = 1 
    
	# 打开相机
    devices = []
    for device in range(0, 5):
        cap = cv2.VideoCapture(device, cv2.CAP_DSHOW)
        if grabbed := cap.read():
            devices.append(device)
        cap.release()
            

    if len(devices) == 0:
        print('设备数据为 0, 退出')
        sys.exit(0)
    
    fps = cap.get(cv2.CAP_PROP_FPS)
	
	# 设置相机亮度
    cap.set(cv2.CAP_PROP_BRIGHTNESS, 150)
    print(cap.get(cv2.CAP_PROP_BRIGHTNESS))
	
	# # 设置相机增益
    cap.set(cv2.CAP_PROP_GAIN, 1)
    print(cap.get(cv2.CAP_PROP_GAIN))
    
    # # 设置相机分辨率
    width, height = Camera_PROP_1280x720
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, width)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, height)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
    # # 初始化opencv显示窗口‘capture’
    cv2.namedWindow('capture', cv2.WINDOW_NORMAL)
    cv2.resizeWindow('capture', width // 2, height // 2)
	
	# 设置相机帧率
    # cap.set(cv2.CAP_PROP_FPS, 30)
    # print(cap.get(cv2.CAP_PROP_FPS))
	
	# 新建一个'test.avi'文件
    # fps = 8
    # fourcc = cv2.VideoWriter_fourcc(*'XVID')
    # out = cv2.VideoWriter('test.avi', fourcc, fps, (width, height))
	
	# 将相机录制到'test.avi'文件中
    while (int(cv2.getWindowProperty('capture', 0)) != -1):
        ret, frame = cap.read() # 获取当前帧
        if not ret:
            continue
        frame = cv2.flip(frame, 1, dst=None) # 镜像
        # out.write(frame) # 写入帧
        imshow("capture", frame) # 显示
        cv2.waitKey(1)

		# 判断相机退出
        input = cv2.waitKey(1) & 0xFF
        if input == ord('w'):
            cv2.imwrite("%s/%d.png" % (class_name, index), frame)
            print("%s: %d 张图片" % (class_name, index))
            index += 1
        elif input == ord('q'):
            break
	# 释放内存
    # out.release()
    cap.release()
    cv2.destroyAllWindows()
    cv2.waitKey(1)
